import os
import numpy as np
import cv2
import gradio as gr


def get_video_paths(fp):
    video_paths = []
    with open(fp, "r") as f:
        for line in f:
            line = line.strip()
            if len(line) == 0:
                continue
            video_paths.append(line)
    
    return video_paths


def get_a_video(state):
    print("here", state["N"])
    assert state["N"] > 0
    idx = state["idx"]
    video_path = state["video_paths"][idx]
    print(idx, video_path)
    state["idx"] = (idx + 1) % state["N"]
    return state, video_path


if __name__ == "__main__":
    
    ################################################################################
    video_paths_files = [
        "/home/liuyun/HHO-dataset/dataset_visualization/retargeting_video_paths_只有内层优化.txt",
        "/home/liuyun/HHO-dataset/dataset_visualization/retargeting_video_paths_内层+外层(无打分器).txt",
        "/home/liuyun/HHO-dataset/dataset_visualization/retargeting_video_paths_内层+外层+打分器.txt",
    ]
    ################################################################################
    
    with gr.Blocks() as iface:
        
        all_video_paths = []
        for i in range(len(video_paths_files)):
            video_paths = get_video_paths(video_paths_files[i])
            all_video_paths.append(gr.State({
                "idx": 0,
                "N": len(video_paths),
                "video_paths": video_paths,
            }))
        
        gr.Markdown("""<p><h1 align="center">Visualization of Different Retargeting Designs</h1></p>""")
        with gr.Row():
            with gr.Column():
                gr.Markdown("""<p>只有内层优化</p>""")
                video0 = gr.Video(autosize=False, visible=True)
                video0_button = gr.Button(value="Next", interactive=True, variant="primary")
            with gr.Column():
                gr.Markdown("""<p>内层优化+外层优化, 没有打分器</p>""")
                video1 = gr.Video(autosize=False, visible=True)
                video1_button = gr.Button(value="Next", interactive=True, variant="primary")
        with gr.Row():
            with gr.Column():
                gr.Markdown("""<p>内层优化+外层优化+打分器</p>""")
                video2 = gr.Video(autosize=False, visible=True)
                video2_button = gr.Button(value="Next", interactive=True, variant="primary")
            # with gr.Column():
            #     video3 = gr.Video(autosize=False, visible=True)
            #     video3_button = gr.Button(value="Next", interactive=True, variant="primary")
        
        # video0_button.click(fn=get_a_video, inputs=[], outputs=[video0])
        video0_button.click(fn=get_a_video, inputs=[all_video_paths[0]], outputs=[all_video_paths[0], video0])
        video1_button.click(fn=get_a_video, inputs=[all_video_paths[1]], outputs=[all_video_paths[1], video1])
        video2_button.click(fn=get_a_video, inputs=[all_video_paths[2]], outputs=[all_video_paths[2], video2])
        # video3_button.click(fn=get_a_video, inputs=[all_video_paths[3]], outputs=[all_video_paths[3], video3])
            
    iface.queue(concurrency_count=1)
    iface.launch(debug=True, enable_queue=True, server_port=1111, server_name="0.0.0.0", share=True)
